#include "stdafx.h"
#include "volumereader.h"
#include "scoped_array.h"
#include "NTFSBootSector.h"
#include "MFT.h"
#include "Attrs.h"

#include<vector>
#include "Runlist.h"
#include "MemoryBuffer.h"


VolumeReader::VolumeReader(QString volumePath,QObject *parent)
	: QObject(parent),masterMFT(this,-1)
{
	m_buffer = new MemoryBuffer(this);
	
	hVol = CreateFile(volumePath.toStdWString().c_str(),
		GENERIC_READ,
		FILE_SHARE_READ|FILE_SHARE_WRITE,
		NULL,
		OPEN_EXISTING,
		0,
		NULL
		);
	if (hVol==INVALID_HANDLE_VALUE)
		errorClose("Unable to open disk or file");
	else {
		//x Q_ASSERT(hVol!=INVALID_HANDLE_VALUE);
		bootSect = new NTFSBootSector();
		read(0,512,(char*)bootSect);
		SetupFSInfo();
	}
	
}



VolumeReader::~VolumeReader()
{
	SafeCloseHandle(hVol);
	SAFE_DELETE(bootSect);
}

bool VolumeReader::isReady()
{
	 return (hVol != INVALID_HANDLE_VALUE );
}

QString VolumeReader::getError()
{
	return _errStr;
}

const QByteArray VolumeReader::read( UINT64 start,UINT64 len )
{
	bool isOk;	 
	if (!isReady()) {
		return QByteArray((const char*)NULL,0);
	}
	QByteArray ret;
	const int BPS = 512;
	LARGE_INTEGER startLI;
	LARGE_INTEGER toLI;
	DWORD BytesRead;
	char buffer[BPS];
	UINT64 _blkStart = start/BPS*BPS;
	if (start/BPS*BPS!=start ) {//<read not start at sector
		//read the first bytes
		startLI.QuadPart = _blkStart;
		UINT32 _preSize = start - _blkStart;
		bool seekRet = SetFilePointerEx(hVol,startLI,&toLI,FILE_BEGIN);
		isOk = ReadFile(hVol,buffer,BPS,&BytesRead,NULL);
		ret.append(buffer+_preSize,min(BPS-_preSize,len));//<dont read more than len
		_blkStart+=BPS;
	}
	startLI.QuadPart = _blkStart;
	bool seekRet = SetFilePointerEx(hVol,startLI,&toLI,FILE_BEGIN);
	Q_ASSERT(seekRet);
	while (_blkStart<start+len) {
		isOk = ReadFile(hVol,buffer,BPS,&BytesRead,NULL);
		int lenToAppend = BPS;
		if (_blkStart+BPS>=start+len) 
			lenToAppend = start+len-_blkStart;
		ret.append(buffer,lenToAppend);
		_blkStart+=BPS;
	}
 	
	
	Q_ASSERT(ret.size()==len);
	return ret;
}

//************************************
// Method:    read
// FullName:  VolumeReader::read
// Access:    public 
// Returns:    void
// Qualifier:
// Parameter: UINT64 start : relative to the start of volume
// Parameter: UINT64 len
// Parameter: char * target
//************************************
 void VolumeReader::read( UINT64 start,UINT64 len,char * target ) const
{	
	if (m_buffer->has(start,len)) {
		memcpy(target,m_buffer->getData(start,len),len);
		return;
	} else {
#ifdef _DEBUG
		//fprintf(stderr,QString("buffer miss start %1 len %2\n").arg(start).arg(len).toStdString().c_str());
#endif
	}
	static const int BPS = 512;
	LARGE_INTEGER startLI;
	LARGE_INTEGER toLI;
	DWORD BytesRead=0;
	UINT64 cursor = 0;
	UINT64 _blkStart = start/BPS*BPS;//start pos of reading
	if (start/BPS*BPS!=start ) {//<read not start at sector
		//read the first bytes
		startLI.QuadPart = _blkStart;
		UINT32 _preSize = start - _blkStart;
		bool bRet= SetFilePointerEx(hVol,startLI,&toLI,FILE_BEGIN);
		Q_ASSERT(bRet);
		char buffer[BPS];
		bRet = ReadFile(hVol,buffer,BPS,&BytesRead,NULL);
		Q_ASSERT(bRet);
		//ret.append(buffer+_preSize,min(BPS-_preSize,len));
		//<dont read more than len
		memcpy(target,buffer+_preSize,min(BPS-_preSize,len));
		_blkStart+=BPS;//then reading pos advances
		cursor+=min(BPS-_preSize,len);
	}
	Q_ASSERT(cursor<=len);
	if (cursor==len)//read complete
		return;
	startLI.QuadPart = _blkStart;
	bool seekRet = SetFilePointerEx(hVol,startLI,&toLI,FILE_BEGIN);
	Q_ASSERT(seekRet);
	UINT64 _blkEnd = (start+len)/BPS*BPS;//end pos of sectors;
	if (_blkStart!=_blkEnd) {
		//read the whole sectors
		Q_ASSERT(_blkEnd-_blkStart<=1024);
		char buffer[1024];
		bool bRet = ReadFile(hVol,buffer,_blkEnd-_blkStart,&BytesRead,NULL);
		Q_ASSERT(bRet && BytesRead==_blkEnd-_blkStart);
		memcpy(target+cursor,buffer,_blkEnd-_blkStart);
		cursor+=_blkEnd-_blkStart;
	}
	//
	//read tails
	if (_blkEnd<start+len) {
		char buffer[BPS];
		bool bRet = ReadFile(hVol,buffer,BPS,&BytesRead,NULL);
		Q_ASSERT(bRet );
		Q_ASSERT(BytesRead==BPS);
		memcpy(target+cursor,buffer,start+len-_blkEnd);
		cursor+=start+len-_blkEnd;
	}
	Q_ASSERT(cursor==len);
	m_buffer->add(start,len,target);
	return;
}

void VolumeReader::SetupFSInfo()
{
// 	QByteArray mbr = this->read(0,512);
// 	memcpy(bootSect,mbr.data(),512);
	
	masterMFT = MFTAt(C2Byte(bootSect->MFTStartCluster));
	masterMFT.diskOffSet = C2Byte(bootSect->MFTStartCluster);
	return;
}

void VolumeReader::errorClose( QString errStr, QString file , int line )
{
	_errStr = errStr;
	_errFile = file;
	_errLine = line;
	SafeCloseHandle(hVol);
}

static void fixupMFT( char*mftDat ) 
{
	MFT * header = (MFT*)mftDat;
	for(int i = 1;i<header->FixUpCount;i++) {
		mftDat[512*i-2] = mftDat[header->FixUpOffset+i*2];
		mftDat[512*i-1] = mftDat[header->FixUpOffset+i*2+1];
	}
}


//************************************ 
// Returns:   The actural CMFT constructor 
//************************************
CMFT VolumeReader::MFTAt( UINT64 StartByte )
{
	//UINT64 entryNo = getEntryNo(StartByte);
	CMFT ret(this,StartByte);
	Q_ASSERT(StartByte % MFTSIZE==0);
	QList<ResidentAttrHdr*> resAttrs;
	QList<NonResidentAttrHdr*> nresAttrs;
	QList<RunList> runLists;
	//TODO:change to static 
	boost::shared_array<char> mftData(new char[MFTSIZE]);
	read(StartByte,MFTSIZE,mftData.get());
	fixupMFT(mftData.get());
	MFT *header = (MFT*) mftData.get();
	char * attrPtr = (char*)(mftData.get())+header->FirstAttrOffset;
	while ((*(UINT32*)attrPtr) != 0xffffffff) {
		AttrHeader *newAtt = (AttrHeader*)attrPtr;
		if (newAtt->IsNonResident) {
			//runlist and non-resident attr parsed together
			nresAttrs.push_back((NonResidentAttrHdr*)newAtt);
			runLists.push_back(RunList((NonResidentAttrHdr*)newAtt,bootSect->SectorPerCluster*512));
		} else  {
			resAttrs.push_back((ResidentAttrHdr*)newAtt);
			ResidentAttrHdr * attrListHdr = (ResidentAttrHdr*)newAtt;
			if (newAtt->TypeID==AttrType::Data) {
				ret.recycle = (RecycleSupport*)
					((char*)newAtt+attrListHdr->ContentOffset);
			}
			//if it is attribute list, then special treatment
			if (newAtt->TypeID==AttrType::AttributeList) {
				//TODO:attribute list
	/*			int entryOffset = 0;
				AttrListAttr * listEntry=(AttrListAttr*)(attrPtr+attrListHdr->ContentOffset);
				UINT64 myEntryNo = getEntryNo(StartByte);
				UINT64 myFileRef = ((UINT64)header->SeqNo<<48) | myEntryNo ;
				while (entryOffset<attrListHdr->ContentSize) {
					if (listEntry->AttrLocationRef!=myFileRef) {
						UINT64 locEntryNo = listEntry->AttrLocationRef & (((UINT64)1<<48)-1);//lower 48 bits;
						//append more mft data
						CMFT toAppend = MFTEntry(locEntryNo);
						if (toAppend.getBaseRef()==myFileRef)
							ret.append(toAppend);
					}
					//advance pointer and offsetter
					entryOffset+=listEntry->EntryLength;
					listEntry = (AttrListAttr *)((char*)listEntry+listEntry->EntryLength);*/
			}
		}
		attrPtr+=newAtt->Lenth;
	}
	//set return value
	ret.mftData.push_back(mftData);
	//ret.resAttrs.insert(ret.resAttrs.end(),resAttrs.begin(),resAttrs.end());
	ret.resAttrs.append(resAttrs);
	for (int i =0;i<nresAttrs.size();i++) {
		NonResidentAttrHdr* nres = nresAttrs[i];
		if (ret.hasAttr(nres->Header.TypeID,nres->Header.ID)) {
			ret.combineRunList(nres->Header.TypeID,nres->Header.ID,runLists[i]);
		} else {
			ret.nresAttrs.push_back(nresAttrs[i]);
			ret.runLists.push_back(runLists[i]);
		}
 	}
	return ret;
}

CMFT VolumeReader::MFTEntry( int entry )
{

	UINT64 startByte = masterMFT.getAttPos(AttrType::Data,-1,entry*MFTSIZE);//<startbyte of the #entry MFT
	CMFT ret = MFTAt(startByte);
	return ret;
}

UINT64 VolumeReader::volumeSize()
{
	return bootSect->TotalSectors*bootSect->BytesPerSector;
	
}

int BinCmp( char * a, char* b, int len ) 
{
	for (int i = 0;i<len;i++)
		if (a[i]!=b[i])
			return a[i]-b[i];
	return 0;
}

bool VolumeReader::isMFT( UINT64 startByte )
{
	MFT mft;
	static char mftSignature[] = {"FILE"};
	Q_ASSERT(startByte % MFTSIZE==0);
	read(startByte,sizeof(mft),(char*)&mft);
	//MFT's size is 42, so no need to fixup here
	//fixupMFT((char*)&mft);
	bool ret = true;
	//ret = ret && mft.Flags<=3;//only 0x1 and 0x2 is valid,but $Secure have a new flag 0x8;
	ret  = ret && (BinCmp(mft.Signature,mftSignature,4)==0);//HACK:signature match (not true under XP)
	ret = ret && mft.FixUpCount==3;//3 fixup value that fixups 1024 bytes
	//all offset in mft must < 1024
	ret = ret && (mft.FixUpOffset <1024);
	ret = ret && (mft.FirstAttrOffset<1024);
	ret = ret &&  (mft.UsedMFTBytes<=1024);
	ret = ret && mft.AllocMTFBytes<=1024;
	return ret;
}

UINT64 VolumeReader::clusterSize()
{
	return bootSect->BytesPerSector * bootSect->SectorPerCluster;
}

UINT64 VolumeReader::getEntryNo( UINT64 diskOffSet )
{
	UINT64 logicOffset = masterMFT.runLists[0].PhyToLogic(diskOffSet);
	if ((INT64)logicOffset>=0) {
		Q_ASSERT(logicOffset%1024==0);
		return logicOffset/1024;
	} return -1;

}


